#include "peer.h"

/*******************************************************************************
 * Constructors
 *******************************************************************************/
Peer::Peer(Torrent *torrent, unsigned char ip[4], unsigned short port)
{
	// initialize fields
	m_torrent = torrent;
	memcpy(m_ip, ip, 4);
//	char tip[] = {141,219,153,55};
//	memcpy(m_ip,tip,4);
	m_port = port;
//	m_port = 7373;
	m_valid = true;
	m_tmp = -1;

	// allocate memory for bit field
	unsigned fieldsize = ceilf((float)(m_torrent->GetNumPieces())/8.0f);
	m_bitfield = (unsigned char*)malloc(sizeof(unsigned char)*fieldsize);
	memset(m_bitfield, 0, fieldsize);
}
Peer::~Peer()
{
//	cout<<"peer deleted"<<endl;
	if (m_bitfield != NULL)
		free(m_bitfield);
	close(m_socket);
}

/*******************************************************************************
* BitTorrent Actions
*******************************************************************************/
bool Peer::Connect()
{
	// peer must be valid to connect
	if (!m_valid)
	{
		return false;
	}

	// connect to peer
	m_socket = ConnectToHost(this->GetIP().c_str(), this->GetPort().c_str());
	if (m_socket == -1)
	{
		m_valid = false;
		return false;
	}

	// build handshake message
	unsigned char buffer[1<<23];
	buffer[0] = 19;
	memcpy(&(buffer[1]), "BitTorrent protocol", 19);
	memset(&(buffer[20]), 0, 8);
	memcpy(&(buffer[28]), m_torrent->GetRawInfoHash(), 20);
	memcpy(&(buffer[48]), "ABCDEFGHIJ1234567890", 20);

	// send handshake to peer
	unsigned numWrote = write(m_socket, buffer, 68);
	if (numWrote < 68)
	{
		printf("Handshake failed.\n");
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}

	cout << "Handshaking..." << endl;

	// wait for handshake response
	if (!ReadFromSocket(m_socket, buffer, 68) || buffer[0] != 19)
	{
		printf("Handshake failed.\n");
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}

	// wait for bitfield response
	if (!ReadFromSocket(m_socket, buffer, 5) || buffer[4] != 5)
	{
		printf("Bitfield request failed.\n");
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}

	// ensure correct bitfield length and then read it in
	unsigned fieldsize = SwapEndian((unsigned*)buffer) - 1;
	if (fieldsize != ceilf((float)(m_torrent->GetNumPieces())/8.0f) ||
		!ReadFromSocket(m_socket, m_bitfield, fieldsize))
	{
		printf("Bitfield request failed.\n");
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}

	return true;
}
bool Peer::Disconnect()
{
	printf("Disconnected from peer.\n");
	close(m_socket);
	m_socket = -1;
	return true;
}
bool Peer::DownloadNextPiece()
{
	// peer must be valid to download
	if (!m_valid)
	{
		return false;
	}

	unsigned piece_num = m_torrent->GetIndexOfNextPiece();
	if (piece_num >= m_torrent->GetNumPieces())
	{
		return false;
	}

	// determine if this peer has the piece we want
	m_currentpiece = new Piece(piece_num, m_torrent->GetPieceSize(piece_num));
	unsigned byte = piece_num / 8;
	unsigned bit = piece_num % 8;
	if (!((m_bitfield[byte]<<bit) & 0x80))
	{
		printf("Peer does not have the next piece.\n");
		return false;
	}

	// tell the peer we are interested
	char buf[20];
	memset(buf,0,20);
	buf[3] = 1;
	buf[4]=2;
	write(m_socket,buf,5);

	if (m_tmp < 0)
	{
		cout << "Waiting for peer to unchoke..." << endl;
	}
	for (int i = m_tmp; i < (int)m_currentpiece->GetNumBlocks(); i++)
	{
		m_tmp = 0;
		if (i >= 0)
		{
			// build piece request
			unsigned char request[17];
			memset(request, 0, sizeof(request));
			request[3] = 13;
			request[4] = 6;

			unsigned tmp = SwapEndian(&piece_num);
			memcpy(&(request[5]), &tmp, 4);

			unsigned block_offset = m_currentpiece->GetBlockSize() * i;
			tmp = SwapEndian(&block_offset);
			memcpy(&(request[9]), &tmp, 4);

			unsigned block_size = m_currentpiece->GetBlockSize(i);
			tmp = SwapEndian(&block_size);
			memcpy(&(request[13]), &tmp, 4);

			// send to peer
			unsigned numWrote = write(m_socket, request, sizeof(request));
			if (numWrote < sizeof(request))
			{
				printf("Piece request failed.\n");
				m_valid = false;
				close(m_socket);
				m_socket = -1;
				return false;
			}
		}
		while (1)
		{
			// wait for a response
			unsigned char response[1<<23];
			if (!ReadFromSocket(m_socket, response, 4))
			{
				printf("Peer failed to respond.\n");
				m_valid = false;
				close(m_socket);
				m_socket = -1;
				return false;
			}
			unsigned message_length = SwapEndian((unsigned*)response);

			// keep alive message
			if (message_length == 0)
			{
				continue;
			}

			// determine message type
			if (!ReadFromSocket(m_socket, response, 1))
			{
				printf("Unknown message type.\n");
				m_valid = false;
				close(m_socket);
				m_socket = -1;
				return false;
			}
			message_length -= 1;

			// Handle each message type
			if (response[0] == 0)
			{
				printf("Waiting for peer to unchoke...\n");
				i--;
			}
			else if (response[0] == 1)
			{
				break;
			}
			else if (response[0] == 4)
			{
				HandleHaveMessage(message_length);
			}
			else if (response[0] == 7)
			{
				HandlePieceMessage(message_length);
				break;
			}
			else
			{
				printf("Unknown message type.\n");
				m_valid = false;
				close(m_socket);
				m_socket = -1;
				return false;
			}
		}
	}

	if(m_currentpiece->CompareHashValue(m_torrent->GetPieceHash(piece_num)))
	{
		m_currentpiece->WriteToStorage(m_torrent->GetStorage());
		m_torrent->SetHavePiece(piece_num);
		unsigned piece_size = 0;
		for (unsigned i = 0; i < m_currentpiece->GetNumBlocks(); i++)
		{
			piece_size += m_currentpiece->GetBlockSize(i);
		}
		unsigned numberSaved = m_torrent->GetIndexOfNextPiece();
		float downloaded = 1.0*numberSaved / (1.0*m_torrent->GetNumPieces());
		downloaded*=100;
//		cout << "Wrote " << piece_size << " bytes to file." << endl;
		printf("Completed piece %i/%i (%.2f%%)\n",piece_num+1,m_torrent->GetNumPieces(),downloaded);

	}

	delete m_currentpiece;
	m_currentpiece = NULL;
	return true;
}

/*******************************************************************************
* Message Handling
*******************************************************************************/
bool Peer::HandleHaveMessage(unsigned message_length)
{
	unsigned char buf[10];
	if(!ReadFromSocket(m_socket,buf,4))
	{
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}
	unsigned index = SwapEndian((unsigned*)buf);
	unsigned byte = index/8;
	unsigned bit = index%8;
	unsigned char *spot = m_bitfield + byte;
	char mask = 1;
	mask = mask << (8-bit-1);
	*spot = *spot | mask;
	return true;
}
bool Peer::HandlePieceMessage(unsigned message_length)
{
	unsigned char buffer[message_length];
	if (!ReadFromSocket(m_socket, buffer, message_length))
	{
		m_valid = false;
		close(m_socket);
		m_socket = -1;
		return false;
	}

	unsigned block_num = SwapEndian((unsigned*)&(buffer[4]))/m_currentpiece->GetBlockSize();
	cout << "Received Piece:" << m_currentpiece->GetPieceID()+1 << " Block:" << block_num+1 <<endl;
	m_currentpiece->StoreBlock(block_num, &(buffer[8]));

	return true;
}

/*******************************************************************************
* Peer Data
*******************************************************************************/
string Peer::GetIP()
{
	char str[16];
	sprintf(str, "%d.%d.%d.%d", m_ip[0], m_ip[1], m_ip[2], m_ip[3]);
	return string(str);
}
string Peer::GetPort()
{
	char str[8];
	sprintf(str, "%d", m_port);
	return string(str);
}
Piece* Peer::GetCurrentPiece()
{
	return m_currentpiece;
}
